#include "stdafx.h"
#include "WSJ.UTIL.h"
#define ARRSIZE(x)	(sizeof(x)/sizeof(x[0]))

/////////////////////////////
//Helper Window class
//////////////////////////////////////////////////////////////////////

CAsyncSocketEx::CAsyncSocketEx()
{
	m_SocketData.hSocket = INVALID_SOCKET;
	m_SocketData.nSocketIndex = -1;
	m_pLocalAsyncSocketExThreadData = 0;

	m_pAsyncGetHostByNameBuffer = NULL;
	m_hAsyncGetHostByNameHandle = NULL;
}
CAsyncSocketEx::~CAsyncSocketEx()
{
	Close();
	FreeAsyncSocketExInstance();
}

BOOL CAsyncSocketEx::Create(UINT nSocketPort /*=0*/, int nSocketType /*=SOCK_STREAM*/, 
							long lEvent /*=FD_READ | FD_WRITE | FD_OOB | FD_ACCEPT | FD_CONNECT | FD_CLOSE*/, 
							LPCSTR lpszSocketAddress /*=NULL*/, BOOL bReuseAddr /*=FALSE*/ )
{
	//Close the socket, although this should not happen
	if (GetSocketHandle() != INVALID_SOCKET) {
		ASSERT(0);
		WSASetLastError(WSAEALREADY);
		return FALSE;
	}

	if (!InitAsyncSocketExInstance()) {
		ASSERT(0);
		WSASetLastError(WSANOTINITIALISED);
		return FALSE;
	}
	SOCKET hSocket = socket(AF_INET, nSocketType, 0);
	if (hSocket == INVALID_SOCKET)
		return FALSE;
	m_SocketData.hSocket = hSocket;

	AttachHandle(hSocket);
	if (!AsyncSelect(lEvent)) {
		Close();
		return FALSE;
	}

	if (bReuseAddr)
	{
		// Has to be done right before binding the socket!
		int iOptVal = 1;
		VERIFY( SetSockOpt(SO_REUSEADDR, &iOptVal, sizeof iOptVal) );
	}

	if (!Bind(nSocketPort, lpszSocketAddress)) {
		Close();
		return FALSE;
	}

	return TRUE;
}

void CAsyncSocketEx::OnReceive(int /*nErrorCode*/)
{
}

void CAsyncSocketEx::OnSend(int /*nErrorCode*/)
{
}

void CAsyncSocketEx::OnConnect(int /*nErrorCode*/)
{
}

void CAsyncSocketEx::OnAccept(int /*nErrorCode*/)
{
}

void CAsyncSocketEx::OnClose(int /*nErrorCode*/)
{
}

BOOL CAsyncSocketEx::OnHostNameResolved(const SOCKADDR_IN * /*pSockAddr*/)
{
	return TRUE;
}

BOOL CAsyncSocketEx::Bind(UINT nSocketPort, LPCSTR lpszSocketAddress)
{
	SOCKADDR_IN sockAddr = {0};
	if (lpszSocketAddress == NULL)
		sockAddr.sin_addr.s_addr = htonl(INADDR_ANY);
	else
	{
		sockAddr.sin_addr.s_addr = inet_addr(lpszSocketAddress);
		if (sockAddr.sin_addr.s_addr == INADDR_NONE) {
			WSASetLastError(WSAEINVAL);
			return FALSE;
		}
	}
	sockAddr.sin_family = AF_INET;
	sockAddr.sin_port = htons((u_short)nSocketPort);
	return Bind((SOCKADDR*)&sockAddr, sizeof(sockAddr));
}

BOOL CAsyncSocketEx::Bind(const SOCKADDR* lpSockAddr, int nSockAddrLen)
{
	if (!bind(m_SocketData.hSocket, lpSockAddr, nSockAddrLen))
		return TRUE;
	else
		return FALSE;
}

void CAsyncSocketEx::AttachHandle(SOCKET /*hSocket*/)
{
	ASSERT( m_pLocalAsyncSocketExThreadData );
	VERIFY( m_pLocalAsyncSocketExThreadData->m_pHelperWindow->AddSocket(this, m_SocketData.nSocketIndex) );
}

void CAsyncSocketEx::DetachHandle(SOCKET /*hSocket*/)
{
	ASSERT( m_pLocalAsyncSocketExThreadData );
	if (!m_pLocalAsyncSocketExThreadData)
		return;

	ASSERT( m_pLocalAsyncSocketExThreadData->m_pHelperWindow );
	if (!m_pLocalAsyncSocketExThreadData->m_pHelperWindow)
		return;

	VERIFY( m_pLocalAsyncSocketExThreadData->m_pHelperWindow->RemoveSocket(this, m_SocketData.nSocketIndex) );
}

void CAsyncSocketEx::Close()
{
	if (m_SocketData.hSocket != INVALID_SOCKET)
	{
		VERIFY( closesocket(m_SocketData.hSocket) != SOCKET_ERROR );
		DetachHandle(m_SocketData.hSocket);
		m_SocketData.hSocket = INVALID_SOCKET;
	}
	delete[] m_pAsyncGetHostByNameBuffer;
	m_pAsyncGetHostByNameBuffer = NULL;
	if (m_hAsyncGetHostByNameHandle) {
		WSACancelAsyncRequest(m_hAsyncGetHostByNameHandle);
		m_hAsyncGetHostByNameHandle = NULL;
	}
}

BOOL CAsyncSocketEx::InitAsyncSocketExInstance()
{
	//Check if already initialized
	if (m_pLocalAsyncSocketExThreadData)
		return TRUE;

	DWORD id = GetCurrentThreadId();
	BOOL bResult = TRUE;
	m_sGlobalCriticalSection.Lock();

	try{
		//Get thread specific data
		if (m_spAsyncSocketExThreadDataList)
		{
			t_AsyncSocketExThreadDataList *pList = m_spAsyncSocketExThreadDataList;
			while (pList)
			{
				ASSERT( pList->pThreadData );
				ASSERT( pList->pThreadData->nInstanceCount > 0 );

				if (pList->pThreadData->nThreadId == id)
				{
					m_pLocalAsyncSocketExThreadData = pList->pThreadData;
					m_pLocalAsyncSocketExThreadData->nInstanceCount++;
					break;
				}
				pList = pList->pNext;
			}
			//Current thread yet has no sockets
			if (!pList)
			{
				//Initialize data for current thread
				pList = new t_AsyncSocketExThreadDataList;
				pList->pNext = m_spAsyncSocketExThreadDataList;
				m_spAsyncSocketExThreadDataList = pList;
				m_pLocalAsyncSocketExThreadData = new t_AsyncSocketExThreadData;
				m_pLocalAsyncSocketExThreadData->nInstanceCount = 1;
				m_pLocalAsyncSocketExThreadData->nThreadId = id;
				m_pLocalAsyncSocketExThreadData->m_pHelperWindow = new CAsyncSocketExHelperWindow;
				m_spAsyncSocketExThreadDataList->pThreadData = m_pLocalAsyncSocketExThreadData;
			}
		}
		else
		{	//No thread has instances of CAsyncSocketEx; Initialize data
			m_spAsyncSocketExThreadDataList = new t_AsyncSocketExThreadDataList;
			m_spAsyncSocketExThreadDataList->pNext = 0;
			m_pLocalAsyncSocketExThreadData = new t_AsyncSocketExThreadData;
			m_pLocalAsyncSocketExThreadData->nInstanceCount = 1;
			m_pLocalAsyncSocketExThreadData->nThreadId = id;
			m_pLocalAsyncSocketExThreadData->m_pHelperWindow = new CAsyncSocketExHelperWindow;
			m_spAsyncSocketExThreadDataList->pThreadData = m_pLocalAsyncSocketExThreadData;
		}
	}
	catch(...){
		TRACE("Unknown exception in CAsyncSocketEx::InitAsyncSocketExInstance()\n");
		ASSERT(0);
		bResult = FALSE;
	}

	m_sGlobalCriticalSection.Unlock();

	return bResult;
}

void CAsyncSocketEx::FreeAsyncSocketExInstance()
{
	//Check if already freed
	if (!m_pLocalAsyncSocketExThreadData)
		return;

	DWORD id = m_pLocalAsyncSocketExThreadData->nThreadId;
	m_sGlobalCriticalSection.Lock();

	try{
		ASSERT( m_spAsyncSocketExThreadDataList );
		t_AsyncSocketExThreadDataList *pList = m_spAsyncSocketExThreadDataList;
		t_AsyncSocketExThreadDataList *pPrev = 0;

		//Search for data for current thread and decrease instance count
		while (pList)
		{
			ASSERT( pList->pThreadData );
			ASSERT( pList->pThreadData->nInstanceCount > 0 );

			if (pList->pThreadData->nThreadId == id)
			{
				ASSERT( m_pLocalAsyncSocketExThreadData == pList->pThreadData );
				m_pLocalAsyncSocketExThreadData->nInstanceCount--;

				//Freeing last instance?
				//If so, destroy helper window
				if (!m_pLocalAsyncSocketExThreadData->nInstanceCount)
				{
					delete m_pLocalAsyncSocketExThreadData->m_pHelperWindow;
					delete m_pLocalAsyncSocketExThreadData;
					if (pPrev)
						pPrev->pNext = pList->pNext;
					else
						m_spAsyncSocketExThreadDataList = pList->pNext;
					delete pList;
					break;
				}
				break;
			}
			pPrev = pList;
			pList = pList->pNext;
			ASSERT( pList );
		}
	}
	catch(...){
		TRACE("Unknown exception in CAsyncSocketEx::FreeAsyncSocketExInstance()\n");
		ASSERT(0);
	}

	m_sGlobalCriticalSection.Unlock();
}

int CAsyncSocketEx::Recv(void* lpBuf, int nBufLen, int nFlags /*=0*/)
{
	return recv(m_SocketData.hSocket, (LPSTR)lpBuf, nBufLen, nFlags);
}

int CAsyncSocketEx::Send(const void* lpBuf, int nBufLen, int nFlags /*=0*/)
{
	return send(m_SocketData.hSocket, (LPSTR)lpBuf, nBufLen, nFlags);
}

BOOL CAsyncSocketEx::Connect(LPCSTR lpszHostAddress, UINT nHostPort)
{
	ASSERT( lpszHostAddress != NULL );

	{
		SOCKADDR_IN sockAddr = {0};
		sockAddr.sin_addr.s_addr = inet_addr(lpszHostAddress);
		if (sockAddr.sin_addr.s_addr == INADDR_NONE)
		{
			m_pAsyncGetHostByNameBuffer = new char[MAXGETHOSTSTRUCT];
			m_nAsyncGetHostByNamePort = nHostPort;
			m_hAsyncGetHostByNameHandle = WSAAsyncGetHostByName(GetHelperWindowHandle(), WM_SOCKETEX_GETHOST, lpszHostAddress, m_pAsyncGetHostByNameBuffer, MAXGETHOSTSTRUCT);
			if (!m_hAsyncGetHostByNameHandle)
				return FALSE;
			WSASetLastError(WSAEWOULDBLOCK);
			return TRUE;
		}
		sockAddr.sin_family = AF_INET;
		sockAddr.sin_port = htons((u_short)nHostPort);
		return CAsyncSocketEx::Connect((SOCKADDR*)&sockAddr, sizeof(sockAddr));
	}
}

BOOL CAsyncSocketEx::Connect(const SOCKADDR* lpSockAddr, int nSockAddrLen)
{
	return connect(m_SocketData.hSocket, lpSockAddr, nSockAddrLen) != SOCKET_ERROR;

}
BOOL CAsyncSocketEx::GetPeerName(SOCKADDR* lpSockAddr, int* lpSockAddrLen)
{

	if (!getpeername(m_SocketData.hSocket, lpSockAddr, lpSockAddrLen))
		return TRUE;
	else
		return FALSE;
}
BOOL CAsyncSocketEx::GetSockName(SOCKADDR* lpSockAddr, int* lpSockAddrLen)
{
	if (!getsockname(m_SocketData.hSocket, lpSockAddr, lpSockAddrLen))
		return TRUE;
	else
		return FALSE;
}

BOOL CAsyncSocketEx::ShutDown(int nHow /*=sends*/)
{
	{
		if( INVALID_SOCKET==m_SocketData.hSocket )
			return TRUE;

		if (!shutdown(m_SocketData.hSocket, nHow))
			return TRUE;
		else
			return FALSE;
	}
}

SOCKET CAsyncSocketEx::Detach()
{
	SOCKET socket = m_SocketData.hSocket;
	DetachHandle(socket);
	m_SocketData.hSocket = INVALID_SOCKET;
	return socket;
}

BOOL CAsyncSocketEx::Attach(SOCKET hSocket, long lEvent /*= FD_READ | FD_WRITE | FD_OOB | FD_ACCEPT | FD_CONNECT | FD_CLOSE*/)
{
	if (hSocket == INVALID_SOCKET || !hSocket)
		return FALSE;
	VERIFY( InitAsyncSocketExInstance() );
	m_SocketData.hSocket = hSocket;
	AttachHandle(hSocket);
	return AsyncSelect(lEvent);
}

BOOL CAsyncSocketEx::AsyncSelect(long lEvent /*= FD_READ | FD_WRITE | FD_OOB | FD_ACCEPT | FD_CONNECT | FD_CLOSE*/)
{
	ASSERT( m_pLocalAsyncSocketExThreadData );
	m_lEvent = lEvent;
	{
		if( m_SocketData.hSocket==INVALID_SOCKET )
			return TRUE;

		if (!WSAAsyncSelect(m_SocketData.hSocket, GetHelperWindowHandle(), m_SocketData.nSocketIndex + WM_SOCKETEX_NOTIFY, lEvent))
			return TRUE;
		else
			return FALSE;
	}
	return TRUE;
}

BOOL CAsyncSocketEx::Listen(int nConnectionBacklog /*=5*/)
{
	if (!listen(m_SocketData.hSocket, nConnectionBacklog))
		return TRUE;
	else
		return FALSE;
}

BOOL CAsyncSocketEx::Accept(CAsyncSocketEx& rConnectedSocket, SOCKADDR* lpSockAddr /*=NULL*/, int* lpSockAddrLen /*=NULL*/)
{
	ASSERT( rConnectedSocket.m_SocketData.hSocket == INVALID_SOCKET );
	{
		SOCKET hTemp = accept(m_SocketData.hSocket, lpSockAddr, lpSockAddrLen);
		if (hTemp == INVALID_SOCKET)
			return FALSE;
		VERIFY(rConnectedSocket.InitAsyncSocketExInstance());
		rConnectedSocket.m_SocketData.hSocket = hTemp;
		rConnectedSocket.AttachHandle(hTemp);
	}
	return TRUE;
}

BOOL CAsyncSocketEx::IOCtl(long lCommand, DWORD* lpArgument)
{
	return ioctlsocket(m_SocketData.hSocket, lCommand, lpArgument) != SOCKET_ERROR;
}

int CAsyncSocketEx::GetLastError()
{
	return WSAGetLastError();
}

BOOL CAsyncSocketEx::TriggerEvent(long lEvent, long nErrCode)
{
	ASSERT( m_pLocalAsyncSocketExThreadData );
	ASSERT( m_pLocalAsyncSocketExThreadData->m_pHelperWindow );
	ASSERT( m_SocketData.nSocketIndex != -1 );

	//  added by yunchenn, 2006/11/27
	//return PostMessage(GetHelperWindowHandle(), m_SocketData.nSocketIndex + WM_SOCKETEX_NOTIFY, m_SocketData.hSocket, lEvent & 0xFFFF);
	return PostMessage(GetHelperWindowHandle(), m_SocketData.nSocketIndex + WM_SOCKETEX_NOTIFY, m_SocketData.hSocket, (lEvent&0xffff) | (nErrCode<<16));
}

SOCKET CAsyncSocketEx::GetSocketHandle()
{
	return m_SocketData.hSocket;
}

HWND CAsyncSocketEx::GetHelperWindowHandle()
{
	if (!m_pLocalAsyncSocketExThreadData)
		return 0;
	if (!m_pLocalAsyncSocketExThreadData->m_pHelperWindow)
		return 0;
	return m_pLocalAsyncSocketExThreadData->m_pHelperWindow->GetHwnd();
}

BOOL CAsyncSocketEx::GetSockOpt(int nOptionName, void* lpOptionValue, int* lpOptionLen, int nLevel /*=SOL_SOCKET*/)
{
	return getsockopt(m_SocketData.hSocket, nLevel, nOptionName, (LPSTR)lpOptionValue, lpOptionLen) != SOCKET_ERROR;
}

BOOL CAsyncSocketEx::SetSockOpt(int nOptionName, const void* lpOptionValue, int nOptionLen, int nLevel /*=SOL_SOCKET*/)
{
	return setsockopt(m_SocketData.hSocket, nLevel, nOptionName, (LPSTR)lpOptionValue, nOptionLen) != SOCKET_ERROR;
}
void CAsyncSocketEx::OnFailConnect(int /*nErrorCode*/)
{
}
